#pragma once

#include "TlsNative.h"


// TYPEDEFs

typedef HANDLE
(WINAPI * WINAPI_CreateFileW)(
    LPCWSTR /*lpFileName*/,
    DWORD /*dwDesiredAccess*/,
    DWORD /*dwShareMode*/,
    LPSECURITY_ATTRIBUTES /*lpSecurityAttributes*/,
    DWORD /*dwCreationDisposition*/,
    DWORD /*dwFlagsAndAttributes*/,
    HANDLE /*hTemplateFile*/
	);

typedef HANDLE
(WINAPI * WINAPI_CreateFileA)(
    LPCSTR /*lpFileName*/,
    DWORD /*dwDesiredAccess*/,
    DWORD /*dwShareMode*/,
    LPSECURITY_ATTRIBUTES /*lpSecurityAttributes*/,
    DWORD /*dwCreationDisposition*/,
    DWORD /*dwFlagsAndAttributes*/,
    HANDLE /*hTemplateFile*/
	);

typedef BOOL
(WINAPI * WINAPI_WriteFile)(
    HANDLE /*hFile*/,
    LPCVOID /*lpBuffer*/,
    DWORD /*nNumberOfBytesToWrite*/,
    LPDWORD /*lpNumberOfBytesWritten*/,
    LPOVERLAPPED /*lpOverlapped*/
    );

typedef BOOL
(WINAPI * WINAPI_ReadFile)(
    HANDLE /*hFile*/,
    LPVOID /*lpBuffer*/,
    DWORD /*nNumberOfBytesToRead*/,
    LPDWORD /*lpNumberOfBytesRead*/,
    LPOVERLAPPED /*lpOverlapped*/
    );

typedef LPVOID
(WINAPI * WINAPI_VirtualAlloc)(
    LPVOID /*lpAddress*/,
    SIZE_T /*dwSize*/,
    DWORD /*flAllocationType*/,
    DWORD /*flProtect*/
    );

typedef LPVOID
(WINAPI * WINAPI_VirtualAllocEx)(
    HANDLE /*hProcess*/,
    LPVOID /*lpAddress*/,
    SIZE_T /*dwSize*/,
    DWORD /*flAllocationType*/,
    DWORD /*flProtect*/
    );

typedef BOOL
(WINAPI * WINAPI_DuplicateHandle)(
    HANDLE /*hSourceProcessHandle*/,
    HANDLE /*hSourceHandle*/,
    HANDLE /*hTargetProcessHandle*/,
    LPHANDLE /*lpTargetHandle*/,
    DWORD /*dwDesiredAccess*/,
    BOOL /*bInheritHandle*/,
    DWORD /*dwOptions*/
    );

typedef void
(WINAPI * WINAPI_ExitProcess)(
    UINT /*uExitCode*/
    );

typedef HANDLE
(WINAPI * WINAPI_CreateThread)(
    LPSECURITY_ATTRIBUTES /*lpThreadAttributes*/,
    SIZE_T /*dwStackSize*/,
    LPTHREAD_START_ROUTINE /*lpStartAddress*/,
    LPVOID /*lpParameter*/,
    DWORD /*dwCreationFlags*/,
    LPDWORD /*lpThreadId*/
    );

typedef VOID
(WINAPI * WINAPI_ExitThread)(
    DWORD /*dwExitCode*/
    );

typedef BOOL
(WINAPI * WINAPI_TerminateThread)(
    HANDLE /*hThread*/,
    DWORD /*dwExitCode*/
    );

typedef HWND
(WINAPI * WINAPI_CreateWindowExA)(
    DWORD /*dwExStyle*/,
    LPCSTR /*lpClassName*/,
    LPCSTR /*lpWindowName*/,
    DWORD /*dwStyle*/,
    int /*X*/,
    int /*Y*/,
    int /*nWidth*/,
    int /*nHeight*/,
    HWND /*hWndParent*/,
    HMENU /*hMenu*/,
    HINSTANCE /*hInstance*/,
    LPVOID /*lpParam*/
	);

typedef HWND
(WINAPI * WINAPI_CreateWindowExW)(
    DWORD /*dwExStyle*/,
    LPCWSTR /*lpClassName*/,
    LPCWSTR /*lpWindowName*/,
    DWORD /*dwStyle*/,
    int /*X*/,
    int /*Y*/,
    int /*nWidth*/,
    int /*nHeight*/,
    HWND /*hWndParent*/,
    HMENU /*hMenu*/,
    HINSTANCE /*hInstance*/,
    LPVOID /*lpParam*/
	);

typedef HMODULE
(WINAPI * WINAPI_LoadLibraryA)(
    LPCSTR /*lpLibFileName*/
    );

typedef HMODULE
(WINAPI * WINAPI_LoadLibraryW)(
    LPCWSTR /*lpLibFileName*/
    );

typedef HMODULE
(WINAPI * WINAPI_LoadLibraryExA)(
    LPCSTR /*lpLibFileName*/,
    HANDLE /*hFile*/,
    DWORD /*dwFlags*/
    );

typedef HMODULE
(WINAPI * WINAPI_LoadLibraryExW)(
    LPCWSTR /*lpLibFileName*/,
    HANDLE /*hFile*/,
    DWORD /*dwFlags*/
    );

typedef DWORD
(WINAPI * WINAPI_WaitForSingleObject)(
    HANDLE /*hHandle*/,
    DWORD /*dwMilliseconds*/
    );

typedef DWORD
(WINAPI * WINAPI_WaitForMultipleObjects)(
    DWORD /*nCount*/,
    CONST HANDLE * /*lpHandles*/,
    BOOL /*bWaitAll*/,
    DWORD /*dwMilliseconds*/
    );

typedef DWORD
(WINAPI * WINAPI_WaitForSingleObjectEx)(
    HANDLE /*hHandle*/,
    DWORD /*dwMilliseconds*/,
    BOOL /*bAlertable*/
    );

typedef DWORD
(WINAPI * WINAPI_WaitForMultipleObjectsEx)(
    DWORD /*nCount*/,
    CONST HANDLE * /*lpHandles*/,
    BOOL /*bWaitAll*/,
    DWORD /*dwMilliseconds*/,
    BOOL /*bAlertable*/
    );

typedef DWORD
(WINAPI * WINAPI_MsgWaitForMultipleObjects)(
    DWORD /*nCount*/,
    CONST HANDLE * /*pHandles*/,
    BOOL /*fWaitAll*/,
    DWORD /*dwMilliseconds*/,
    DWORD /*dwWakeMask*/);

typedef DWORD
(WINAPI * WINAPI_MsgWaitForMultipleObjectsEx)(
    DWORD /*nCount*/,
    CONST HANDLE * /*pHandles*/,
    DWORD /*dwMilliseconds*/,
    DWORD /*dwWakeMask*/,
    DWORD /*dwFlags*/);

// ORIGINAL FUNCTIONS

/**
 * This class holds static pointers to original implementations
 * of all hooked functions
 */
class OriginalAPI {
public:
	static WINAPI_CreateFileW CreateFileW;
	static WINAPI_CreateFileA CreateFileA;
	static WINAPI_WriteFile WriteFile;
	static WINAPI_ReadFile ReadFile;

	static WINAPI_VirtualAlloc VirtualAlloc;
	static WINAPI_VirtualAllocEx VirtualAllocEx;

	static WINAPI_DuplicateHandle DuplicateHandle;
	static WINAPI_ExitProcess ExitProcess;

	static WINAPI_LoadLibraryA LoadLibraryA;
	static WINAPI_LoadLibraryW LoadLibraryW;
	static WINAPI_LoadLibraryExA LoadLibraryExA;
	static WINAPI_LoadLibraryExW LoadLibraryExW;

	static WINAPI_CreateThread CreateThread;
	static WINAPI_ExitThread ExitThread;
	static WINAPI_TerminateThread TerminateThread;

	static WINAPI_CreateWindowExA CreateWindowExA;
	static WINAPI_CreateWindowExW CreateWindowExW;

	static WINAPI_WaitForSingleObject WaitForSingleObject;
	static WINAPI_WaitForSingleObjectEx WaitForSingleObjectEx;
	static WINAPI_WaitForMultipleObjects WaitForMultipleObjects;
	static WINAPI_WaitForMultipleObjectsEx WaitForMultipleObjectsEx;
	static WINAPI_MsgWaitForMultipleObjects MsgWaitForMultipleObjects;
	static WINAPI_MsgWaitForMultipleObjectsEx MsgWaitForMultipleObjectsEx;
};


// DECLARATIONS

HANDLE
WINAPI
HookCreateFileW(
    LPCWSTR lpFileName,
    DWORD dwDesiredAccess,
    DWORD dwShareMode,
    LPSECURITY_ATTRIBUTES lpSecurityAttributes,
    DWORD dwCreationDisposition,
    DWORD dwFlagsAndAttributes,
    HANDLE hTemplateFile
	);

HANDLE
WINAPI
HookCreateFileA(
    LPCSTR lpFileName,
    DWORD dwDesiredAccess,
    DWORD dwShareMode,
    LPSECURITY_ATTRIBUTES lpSecurityAttributes,
    DWORD dwCreationDisposition,
    DWORD dwFlagsAndAttributes,
    HANDLE hTemplateFile
	);

BOOL
WINAPI
HookWriteFile(
    HANDLE hFile,
    LPCVOID lpBuffer,
    DWORD nNumberOfBytesToWrite,
    LPDWORD lpNumberOfBytesWritten,
    LPOVERLAPPED lpOverlapped
    );

BOOL
WINAPI
HookReadFile(
    HANDLE hFile,
    LPVOID lpBuffer,
    DWORD nNumberOfBytesToRead,
    LPDWORD lpNumberOfBytesRead,
    LPOVERLAPPED lpOverlapped
    );

LPVOID
WINAPI
HookVirtualAlloc(
    LPVOID lpAddress,
    SIZE_T dwSize,
    DWORD flAllocationType,
    DWORD flProtect
    );

LPVOID
WINAPI
HookVirtualAllocEx(
    HANDLE hProcess,
    LPVOID lpAddress,
    SIZE_T dwSize,
    DWORD flAllocationType,
    DWORD flProtect
    );

BOOL
WINAPI
HookDuplicateHandle(
    HANDLE hSourceProcessHandle,
    HANDLE hSourceHandle,
    HANDLE hTargetProcessHandle,
    LPHANDLE lpTargetHandle,
    DWORD dwDesiredAccess,
    BOOL bInheritHandle,
    DWORD dwOptions
    );

void
WINAPI
HookExitProcess(
	UINT uExitCode
	);

HANDLE
WINAPI
HookCreateThread(
    LPSECURITY_ATTRIBUTES lpThreadAttributes,
    SIZE_T dwStackSize,
    LPTHREAD_START_ROUTINE lpStartAddress,
    LPVOID lpParameter,
    DWORD dwCreationFlags,
    LPDWORD lpThreadId
    );

VOID
WINAPI
HookExitThread(
    DWORD dwExitCode
    );

BOOL
WINAPI
HookTerminateThread(
    HANDLE hThread,
    DWORD dwExitCode
    );

HWND
WINAPI
HookCreateWindowExA(
    DWORD dwExStyle,
    LPCSTR lpClassName,
    LPCSTR lpWindowName,
    DWORD dwStyle,
    int X,
    int Y,
    int nWidth,
    int nHeight,
    HWND hWndParent,
    HMENU hMenu,
    HINSTANCE hInstance,
    LPVOID lpParam);

HWND
WINAPI
HookCreateWindowExW(
    DWORD dwExStyle,
    LPCWSTR lpClassName,
    LPCWSTR lpWindowName,
    DWORD dwStyle,
    int X,
    int Y,
    int nWidth,
    int nHeight,
    HWND hWndParent,
    HMENU hMenu,
    HINSTANCE hInstance,
    LPVOID lpParam);

HMODULE
WINAPI
HookLoadLibraryA(
    LPCSTR lpLibFileName
    );

HMODULE
WINAPI
HookLoadLibraryW(
    LPCWSTR lpLibFileName
    );

HMODULE
WINAPI
HookLoadLibraryExA(
    LPCSTR lpLibFileName,
    HANDLE hFile,
    DWORD dwFlags
    );

HMODULE
WINAPI
HookLoadLibraryExW(
    LPCWSTR lpLibFileName,
    HANDLE hFile,
    DWORD dwFlags
    );

DWORD
WINAPI
HookWaitForSingleObject(
    HANDLE hHandle,
    DWORD dwMilliseconds
    );

DWORD
WINAPI
HookWaitForMultipleObjects(
    DWORD nCount,
    CONST HANDLE *lpHandles,
    BOOL bWaitAll,
    DWORD dwMilliseconds
    );

DWORD
WINAPI
HookWaitForSingleObjectEx(
    HANDLE hHandle,
    DWORD dwMilliseconds,
    BOOL bAlertable
    );

DWORD
WINAPI
HookWaitForMultipleObjectsEx(
    DWORD nCount,
    CONST HANDLE *lpHandles,
    BOOL bWaitAll,
    DWORD dwMilliseconds,
    BOOL bAlertable
    );

DWORD
WINAPI
HookMsgWaitForMultipleObjects(
    DWORD nCount,
    CONST HANDLE *pHandles,
    BOOL fWaitAll,
    DWORD dwMilliseconds,
    DWORD dwWakeMask);

DWORD
WINAPI
HookMsgWaitForMultipleObjectsEx(
    DWORD nCount,
    CONST HANDLE *pHandles,
    DWORD dwMilliseconds,
    DWORD dwWakeMask,
    DWORD dwFlags);


// STATE INDICATOR

extern platform::TlsFlag s_internalModeFlag;

/**
 * This class is used to keep a track for each thread
 * whether it's currently in "Internal" state or not - when
 * in internal, we should register no new events.
 * The flag can't be changed directly by this class, and
 * this must be accomplished using <i>Setter</i> inner class
 */
class StateIndicator {
	static inline bool & GetThreadFlag(bool &if_undef) {
		bool & b_tls_value = s_internalModeFlag();
		return (((&b_tls_value) == NULL) ? if_undef : b_tls_value);
	}

	static inline void SetFlag(void) { bool b; GetThreadFlag(b) = true; }
	static inline void ResetFlag(void) { bool b; GetThreadFlag(b) = false; }

public:
	static inline bool IsInternalState(void) { bool b = true; return GetThreadFlag(b); }

	/**
	 * This class is used to manage the current state
	 * of the thread - when there's a <i>Setter</i>
	 * variable live on the stack, all hooked functions
	 * should call the original function.
	 */
	class Setter {
		bool should_reset_state;

	public:
		inline Setter(void) : should_reset_state(false) {
			if (!StateIndicator::IsInternalState()) {
				should_reset_state = true;
				StateIndicator::SetFlag();
			}
		}

		inline ~Setter(void) {
			if (should_reset_state)
				StateIndicator::ResetFlag();
		}
	};

	friend class Setter;
};


/**
 * It is iportant to have correct value of "LastError"
 * when returning value to callee from hooks. But "LatError"
 * is updated on every system call (such as TlsGetValue), so
 * must make sure to restore it before returning.
 *
 * NOTE: this issue can be a real pain if not handled
 * corectly, it took me a while to understand the problem
 */
class ErrorCodeRestore {
	DWORD dwErr;

public:
	inline ErrorCodeRestore(void) : dwErr(::GetLastError()) {}
	inline ~ErrorCodeRestore(void) { ::SetLastError(dwErr); }
};
